"""
普通渲染器
"""
# pylint: disable=missing-function-docstring,trailing-whitespace,unused-argument
# pylint: disable=consider-using-f-string,no-value-for-parameter
import time
import math
import numpy as np
import taichi as ti

from .gmath import normalized
from .type_enum import PrimitiveType,LightType
from .render import Renderer
from .misc import (bit_blt,copy_2d)
from .simple_ui import (UiWindow,UiTreeItem)

TARGET_FPS = 1200

@ti.data_oriented
class RendererFrame:
    """
    渲染器框架
    """
    def __init__(self,renderer:Renderer,scale=1.0):
        self.resolution = (int(renderer.resolution[0]*scale),int(renderer.resolution[1]*scale))
        self.screen = ti.Vector.field(3, float, renderer.resolution) # 图像的像素场
        self.window = ti.ui.Window("Ray Tracing Renderer",self.resolution)
        self.uiwin = UiWindow(renderer.resolution[0]-320,0,320,renderer.resolution[1])
        self.root = None
        self.bui = False
        self.bspace = False
        self.backspace = False
        self.renderer = renderer
        self.scene = renderer.bounding.scene
        self.last_mouse_pos = None

    def init_ui(self):
        """
        初始化参数控制界面
        """
        self.root = UiTreeItem(isopen=True)
        self.uiwin.add_ui(self.root,0,0)
        render = UiTreeItem("Render",color=(0.5,0.2,0.2,0.3))
        white = (1.,1.,1.,0.5)
        rgb_color = ((0.6,0.1,0.1,0.5),(0.1,0.6,0.1,0.5),(0.1,0.1,0.6,0.5))
        xyz_color = ((0.3,0.3,0.3,0.4),(0.4,0.4,0.4,0.4),(0.5,0.5,0.5,0.4))
        bounding = self.renderer.bounding     
        def color_cb(v,i):
            bounding.environment_color[None][i] = v
            self.renderer.clear()
        for i in range(3):
            c = rgb_color[i]
            render.add_properties("env %s"%(['r','g','b'][i]),
                                  color_cb,
                                  value=bounding.environment_color[None][i],index=i,
                                  label_color=c,
                                  slider_bg=c)
        def fov_cb(v,i):
            self.renderer.fov_tan[None] = math.tan(v/2)
            self.renderer.clear()
        def hdr_cb(v,i):
            self.renderer.hdr = v
            self.renderer.clear()
        render.add_properties("fov",fov_cb,
                                    value=self.renderer.fov,
                                    minv=np.pi/30,maxv=np.pi/4,
                                    label_color=white)
        render.add_check("HDR",hdr_cb,self.renderer.hdr,label_color=white)
        
        material = UiTreeItem("Material",color=(0.2,0.6,0.2,0.3))
        
        def a_cb(v,i):
            bounding.materials[i].a = max(v,0.02)
            self.renderer.clear()
        def m_color_cb(v,index):
            i = index[0]
            j = index[1]
            bounding.materials[i].color[j] = v
            self.renderer.clear()
        def f_cb(v,index):
            i = index[0]
            j = index[1]
            bounding.materials[i].f[j] = v
            self.renderer.clear()
        def metallic_cb(v,i):
            bounding.materials[i].metallic = v
            self.renderer.clear()
        def mu_cb(v,i):
            bounding.materials[i].mu = v
            self.renderer.clear()
        def tex_cb(v,a):
            bounding.materials[a[0]].tex = a[1] if v else -1
            self.renderer.clear()
        def visible_cb(v,a):
            bounding.bounds[a[0]].type = a[1] if v else PrimitiveType.INVISIABLE.value
            self.renderer.clear()

        for i in range(bounding.material_n):
            m = bounding.materials[i]
            mo = bounding.m2n[i]
            s = "%s%s"%(mo[0]+" " if mo[0] else "",mo[1])
            mesh = mo[3].mesh
            if mo[1]=="sphere":
                s += " R=%.1f"%mesh['meta'][0]
            elif mo[1]=="box":
                s += " (%.1f,%.1f,%.1f)"%(mesh['meta'][0],mesh['meta'][1],mesh['meta'][2])
            elif mo[1]=="quad":
                s += " (%.1f,%.1f)"%(mesh['meta'][0],mesh['meta'][1])
            node = UiTreeItem(s,color=(0.1,0.1,0.5,0.3))
            for j in range(3):
                c = rgb_color[j]
                node.add_properties("color %s"%(['r','g','b'][j]),
                                    m_color_cb,
                                    value=m.color[j],index=(i,j),
                                    label_color=c,
                                    slider_bg=c)
            node.add_properties("a",a_cb,value=m.a,index=i,label_color=white)
            for j in range(3):
                c = rgb_color[j]
                node.add_properties("f %s"%(['r','g','b'][j]),
                                    f_cb,
                                    value=m.f[j],index=(i,j),label_color=c,slider_bg=c)
            node.add_properties("metallic",metallic_cb,
                                value=m.metallic,index=i,label_color=white)                    
            node.add_properties("mu",mu_cb,
                                value=m.mu,minv=1.,maxv=2.,index=i,label_color=white)
            node.add_check("texture",tex_cb,m.tex!=-1,label_color=white,index=(i,m.tex))
            bound_i = mo[2]
            btype = bounding.bounds[bound_i].type
            node.add_check("visible",visible_cb,True,label_color=white,index=(mo[2],btype))
            material.add_child(node)

        def light_color_cb(v,a):
            i,j = a
            light = bounding.light[i]
            light.color[j] = v
            self.renderer.clear()

        def light_power_cb(v,a):
            i = a
            light = bounding.light[i]
            light.power = v
            self.renderer.clear()

        def light_p_cb(v,a):
            i,j = a
            light = bounding.light[i]
            light.p[j] = v
            self.renderer.clear()
        
        def light_dir_cb(v,a):
            i,j = a
            light = bounding.light[i]
            light.dir[j] = v
            p = np.array([light.dir[0],light.dir[1],light.dir[2]])
            light.dir = ti.Vector(normalized(p))
            self.renderer.clear()

        light_node = UiTreeItem("Light",color=(0.8,0.8,0.2,0.3))
        for i in range(bounding.light_n):
            light = bounding.light[i]
            light_type = {LightType.POINT_LIGHT.value:"Point light",
                          LightType.DIRECT_LIGHT.value:"Direct light",
                          LightType.SPOT_LIGHT.value:"Spot light",
                          LightType.AREA_LIGHT.value:"Area light",
                          LightType.SKY_LIGHT.value:"Sky light"}
            node = UiTreeItem("%s"%(light_type[light.type]))
            for j in range(3):
                c = rgb_color[j]
                node.add_properties("color %s"%(['r','g','b'][j]),
                                    light_color_cb,
                                    value=light.color[j],index=(i,j),
                                    label_color=c,
                                    slider_bg=c)
            for j in range(3):
                c = xyz_color[j]
                node.add_properties("p %s"%(['x','y','z'][j]),
                                    light_p_cb,
                                    value=light.p[j],index=(i,j),
                                    minv=-100,maxv=100.,
                                    label_color=c,
                                    slider_bg=c)
            for j in range(3):
                c = xyz_color[j]
                node.add_properties("dir %s"%(['x','y','z'][j]),
                                    light_dir_cb,
                                    value=light.dir[j],index=(i,j),
                                    minv=-1,maxv=1.,
                                    label_color=c,
                                    slider_bg=c)                
            node.add_properties("power",light_power_cb,value=light.power,
                                minv=0.,maxv=2*light.power,index=i,
                                label_color=white)
            light_node.add_child(node)

        self.root.add_child(render)
        self.root.add_child(material)
        self.root.add_child(light_node)
        #摄像机
        def phi_cb(v,a):
            if v !=0:
                self.renderer.lens_phi[None] = 0.1/v
            self.renderer.clear()
        def lens_cb(v,a):
            self.renderer.lens_l[None] = v
            self.renderer.clear()              
        camera = UiTreeItem("Camera",color=(0.2,0.8,0.8,0.3))
        phi = 0.1/self.renderer.lens_phi[None] if self.renderer.lens_phi[None]!=0 else 0.
        camera.add_properties("phi",phi_cb,
                              value=phi,
                              minv=0.1,maxv=10.,
                              label_color=white)
        camera.add_properties("lens",lens_cb,
                              value=self.renderer.lens_l[None],
                              minv=1.,maxv=100.,
                              label_color=white)
        self.root.add_child(camera)
        self.uiwin.update()

    def update_camera(self):
        """
        更新摄像机
        """
        step = 1
        ps = 10
        pressed = False
        movedir = [0,0,0]
        if (self.window.is_pressed(ti.ui.SPACE) or self.window.is_pressed(ti.ui.ESCAPE)) and not self.bspace:
            self.bui = not self.bui
            self.bspace = True
        elif not (self.window.is_pressed(ti.ui.SPACE) or self.window.is_pressed(ti.ui.ESCAPE)) and self.bspace:
            self.bspace = False
        
        if self.bui:
            x,y = self.window.get_cursor_pos()
            b = 0
            if self.window.is_pressed(ti.ui.LMB):
                b = 1 
            elif self.window.is_pressed(ti.ui.RMB):
                b = 2
            xx = x * self.renderer.resolution[0]-(self.renderer.resolution[0]-320)
            yy = (1-y) * self.renderer.resolution[1]
            self.uiwin.on_mouse(b,xx,yy)
            return
        
        isbackspace = self.window.is_pressed(ti.ui.BACKSPACE)
        if isbackspace and not self.backspace:
            self.backspace = True
            self.save_screen()
        elif isbackspace and self.backspace:
            self.backspace = False

        for key,d in [('w',(0,-1,0)),('s',(0,1,0)),('a',(-1,0,0)),('d',(1,0,0)),
                      (ti.ui.UP,(0,0,1)),(ti.ui.DOWN,(0,0,-1))]:
            if self.window.is_pressed(key):
                pressed = True
                movedir[0] += step*d[0]
                movedir[1] += step*d[1]
                movedir[2] += step*d[2]
        if self.scene and self.scene.camera:
            self.scene.camera.move_wsad(movedir[0],movedir[1],movedir[2])

        for key,d in [(ti.ui.LEFT,-1),(ti.ui.RIGHT,1)]:
            if self.window.is_pressed(key):
                pressed = True
                if self.scene and self.scene.camera:
                    self.scene.camera.rotate_look_pos(d * np.pi/36.)

        if self.window.is_pressed(ti.ui.LMB):
            pressed = True
            mouse_pos = np.array(self.window.get_cursor_pos())
            if self.last_mouse_pos is not None:
                d_heading,d_pitch = ps * (self.last_mouse_pos - mouse_pos)*np.pi/180
                #print(d_pitch,d_heading)
                if self.scene.camera:
                    self.scene.camera.rotate_fov(-d_heading,d_pitch)
            self.last_mouse_pos = mouse_pos
        else:
            self.last_mouse_pos = None

        return pressed

    def mix_ui_renderer(self):
        """
        将渲染内容和ui混合
        """
        if self.bui:
            copy_2d(self.renderer.fetch_image(),self.screen)
            bit_blt(self.uiwin.pixels,self.screen,self.uiwin.pos[0],self.uiwin.pos[1])
            return self.screen
        else:
            return self.renderer.fetch_image()

    def save_screen(self):
        img = self.mix_ui_renderer()
        filename = "./screen/%d.png"%(np.random.randint(1,1e9))
        ti.tools.imwrite(img,filename)
        print(filename)

    def loop(self):
        """
        主循环
        """
        self.renderer.set_camera(self.scene.camera)
        canvas = self.window.get_canvas()
        self.init_ui()
        spp = 1
        while self.window.running:
            should_reset_framebuffer = False
            if self.update_camera():
                should_reset_framebuffer = True

            if should_reset_framebuffer:
                self.renderer.set_camera(self.scene.camera)
                self.renderer.reset_framebuffer()

            t = time.time()
            for _ in range(spp):
                self.renderer.accumulate()

            canvas.set_image(self.mix_ui_renderer())

            elapsed_time = time.time() - t
            if elapsed_time * TARGET_FPS > 1:
                spp = int(spp / (elapsed_time * TARGET_FPS) - 1)
                spp = max(spp, 1)
            else:
                spp += 1

            self.window.show()
